nice! give it a test drive with the cifar script in the examples folder

@kashif ^

Still not 100% sure this is correct, but initial results seem promising.

 vq :: rec loss: 0.114 | cmt loss: 0.001 | active %: 21.094

fsq :: rec loss: 0.111 | active %: 58.333

Not parameter matched, so losses aren't really representative, but interesting to see the higher codebook usage.

do you have the training curves for each?

edit: removed the images that were here because the test was broken

Probably isn't representative of FSQ performance in general, but it does seem to be working on some level at least ^^

Ok wait, strike that. Reading the paper a bit more closely, the offsets should be flipped relative to what I put there initially. Doing so nets a much higher active %.

Okay, this seems fine for an example. Curious to see if it works in more realistic scenarios too.

@sekstini getting there! we can let @kashif do a review before getting it merged! thanks for beasting through this

Okay, this seems fine for an example. Curious to see if it works in more realistic scenarios too.

Cool to see the high util out of the box :) I assume this is a fairly shallow AE? Probably the loss gap could be a bit smaller but looks like it's WAI! nice

@sekstini can you also kindly add a section in the README and the appropriate bibtex entry as #73 had started?

Okay, this seems fine for an example. Curious to see if it works in more realistic scenarios too.

Cool to see the high util out of the box :) I assume this is a fairly shallow AE? Probably the loss gap could be a bit smaller but looks like it's WAI! nice

Yeah, the test models are both shallow, and tiny at ~10k parameters. Don't expect this to realistically reflect performance at all. This is the FSQ network.

@sekstini can you also kindly add a section in the README and the appropriate bibtex entry as #73 had started?

Will do 👍

Yeah, the test models are both shallow, and tiny at ~10k parameters. Don't expect this to realistically reflect performance at all. This is the FSQ network.

Makes sense! Yeah one thing with FSQ is that because your codebook is fixed, you need some decent number of layers before and after the quantizer. Usually this is the case (eg the VQ-GANs people train are fairly deep).

Anyway, I think the results you see are what we expect! Good stuff.

For the figures, I revisited your README and it seems it is usually mostly the code, so maybe simply using the cube (PDF here) would be enough or already too much. Your call :)

In the README of our upcoming mini-repo, I also added the table, but again, this might be out of place in your README:

source:

|                  | VQ | FSQ |
|------------------|----|-----|
| Quantization     | argmin_c \|\| z-c \|\| | round(f(z)) |
| Gradients        | Straight Through Estimation (STE) | STE |
| Auxiliary Losses | Commitment, codebook, entropy loss, ... | N/A |
| Tricks           | EMA on codebook, codebook splitting, projections, ...| N/A |
| Parameters       | Codebook | N/A |

lgtm! releasing

thank you @sekstini !

@lucidrains Oh, apparently I pointed this as the fsq branch, so you might need to merge it into master

@sekstini yup, no problem, thank you! 🙏

Thanks everyone! I'll add a link to this repo in our next revision

Added a link to the official README now.

Appreciate this quick port and example code. Just thought I would add that checking reconstructions from the FashionMNIST example appear reasonable

However a bit of surprise noticing it fails if switching back to MNIST

Changing levels/seeds didn't help, so perhaps as @fab-jul mentioned it's just a case of needing more layers before/after quantizer for some cases.

@dribnet Interesting. Not sure why we would see this particular failure mode here, but I made a toy example of residual vector quantization where it's working: https://gist.github.com/sekstini/7f089f71d4b975ec8bde37d878b514d0.

@dribnet adding more layers seems testable

LFQ https://arxiv.org/abs/2310.05737 looks simliar?

nevermind, it is slightly different, will be adding

nevermind, it is slightly different, will be adding

It‘s FSQ with levels = 2 plus an entropy maximization loss, IIUC.

yes, FSQ generalizes it, save for the entropy loss. no ablation of the entropy loss, so unsure how necessary it is, but i'll add it. perhaps that will be their contribution

@fab-jul congrats either way! this could be a winner

@lucidrains thanks!

for us, levels = 2 was suboptimal but I can see how an entropy maximization would fix that. I wonder if that also improves FSQ (although ofc the goal of our paper was to get rid of all aux losses haha)

@sekstini you should pair up with an audio researcher and think hard about this (residual fsq), maybe do a tweetstorm or blogpost if you see anything but do not have bandwidth to write a paper. i'm sure many will be thinking along these lines, looking at the new paper

Funny you mention it, as I'm experimenting with a vocoder based on this idea in this very moment ^^
I should definitely improve my tweeting game though.

@sekstini are you seeing good results?

thinking about a soundstream variation with multi-headed LFQ

@lucidrains

@sekstini are you seeing good results?

Not really, but definitely not because of FSQ.

I have exclusively been toying with various parallel encoding schemes, which I'm guessing are difficult for the model to learn, and I suspect residual quantization would work a lot better.

@sekstini ah got it, thanks!

yea, to do residual, i think the codes will need to be scaled down an order of magnitude, or inputs scaled up (cube within cubes, if you can imagine it), but i haven't worked it out. it is probably a 3 month research project. somebody will def end up trying it..

@sekstini yea, if LFQ pans out over at magvit2, i'll do some improvisation here and maybe someone can do the hard experimental work.

@sekstini almost done! #80

@sekstini almost done! #80

Neat! I made some decent progress on my vocoder with "parallel FSQ", but I'd be interested in swapping this in to compare performance. Feel free to tag me when it's done.

ok it is done, integrated in soundstream over here

may be of interest! lucidrains/magvit2-pytorch#4

Hi,
Thanks for the implementation!
@sekstini or @fab-jul Could you perhaps briefly explain what the purpose of the "shift/offset" is?
The reason why I am asking: For levels=2 "shift" becomes infinite. Should it be (1 + eps) instead of (1-eps)?

@mueller-franzes You may find this comment by Fabian interesting.

Should it be (1 + eps) instead of (1-eps)?

Makes sense to me. I copied it directly from the paper, but at the time there was a tan instead of atanh there, so I didn't notice any issues while testing levels = 2.

As a side note, you may want to check out LFQ if you're interested in the levels = 2 case in particular.

@mueller-franzes @sekstini maybe we should just enforce odd levels for now for FSQ until a correction to the paper comes out?

and yea, agreed with checking out LFQ. so many groups seeing success with it

@mueller-franzes @sekstini maybe we should just enforce odd levels for now for FSQ until a correction to the paper comes out?

Other than the asymmetry being weird, levels > 2 seems fine (actually most of my code has been using even values).

I think switching to (1 + eps) or enforcing levels > 2 makes sense.

Thank you both for the super quick response! That's a good tip, I'll have a look at LFQ next.

@sekstini oh you are right, it is a levels == 2 problem

ok let's just go with 1 + eps! thank you both